#include "common.h"
#include "vsprintf.h"

extern u32 __romhax_start;
extern u32 __bss_start;
extern u32 __bss_end;
extern u32 __got_start;
extern u32 __got_end;
extern u32 __romhax_end;

typedef void (*ep_t)(void);

typedef struct {
	char is_usb3;
	char init_hw_done;
	char init_proto_done;
	char unk0;

	int (*init_hw)(void);
	int (*init_proto)(void);

	void *ep1_out;
	void *ep1_out_get_len_proc_ep0;
	void (*ep1_out_imm)(void *buffer, u32 size, u32 *num_xfer);

	void *ep1_in;
	void *ep1_in_get_len_proc_ep0;
	int (*ep1_in_imm)(void *buffer, u32 size, u32 *num_xfer);

	void *ep0_stall;
} rcm_transport_t;

static const rcm_transport_t *rcm_transport = (rcm_transport_t *)0x40003114;

static u32 rom_recvbuf(void *buffer, u32 size) {
	u32 num_xfer;
	rcm_transport->ep1_out_imm(buffer, size, &num_xfer);
	return num_xfer;
}

static u32 rom_sendbuf(void *buffer, u32 size) {
	u32 num_xfer;
	rcm_transport->ep1_in_imm(buffer, size, &num_xfer);
	return num_xfer;
}

static u32 recv_to(void *addr, u32 size)
{
	const u32 chunk_max = 1024;
	u32 left = size, chunk;
	
	while (left > 0) {
		chunk = left;
		if (chunk > chunk_max)
			chunk = chunk_max;

		chunk = rom_recvbuf(addr, chunk);

		addr += chunk;
		left -= chunk;
	}

	return size - left;
}

int printf(const char* fmt, ...) {
	char buffer[512];
	va_list args;
	int i;
	va_start(args, fmt);
	i = vsprintf(buffer, fmt, args);
	va_end(args);
	return rom_sendbuf(buffer, i);
}

#define PMC_BASE		0x7000e400
#define PMC_CNTRL		0x000
#define  PMC_CNTRL_MAIN_RST	(1 << 4)
#define PMC_SCRATCH0		0x050
#define  PMC_SCRATCH0_MODE_RCM	(1 << 1)

static void enterrcm() {
	or32(PMC_BASE + PMC_SCRATCH0, PMC_SCRATCH0_MODE_RCM);
	or32(PMC_BASE + PMC_CNTRL, PMC_CNTRL_MAIN_RST);
}

#define UINT32TOBUF(b, o, val) \
	do { \
		b[o + 0] = (val >> 24) & 0xff; \
		b[o + 1] = (val >> 16) & 0xff; \
		b[o + 2] = (val >> 8) & 0xff; \
		b[o + 3] = val & 0xff; \
	} while (0);

static void _main() {
	/* see coreboot's src/mainboard/nintendo/switch/memlayout.ld */
	const u32 len = 28 * 1024;
	const u32 addr = 0x40010000;
	char buffer[8];
	ep_t ep;

	printf("CBFS\n");

	UINT32TOBUF(buffer, 0, 0)
	UINT32TOBUF(buffer, 4, len)
	rom_sendbuf(buffer, 8);
	recv_to((void *)addr, len);

	ep = (ep_t)addr;
	ep();

	enterrcm();
}

__attribute__((section(".init")))
void relocator_stub() {
	/* Warning: at this point, the GOT is not relocated yet. This means all
	 * the linker script defined symbols (__romhax_start, etc.) are
	 * relative to the current module base. Since this function is
	 * guaranteed to be the first thing emitted, we use its address as the
	 * module base. */

	/* Relocate to end of IRAM. */
	u32 IRAM_END = 0x40040000;
	u32 romhax_size = (((u32)&__romhax_end - (u32)&__romhax_start) + 3) & ~0x3;

	u8* new_base = (u8*)(IRAM_END - romhax_size);
	u8* curr_base = (u8*)relocator_stub - 1; /* Thumb adjustment. */

	ep_t ep;

	memcpy(new_base, curr_base, romhax_size);

	/* Clear BSS. */
	u8* bss_start = new_base + (u32)&__bss_start;
	u8* bss_end = new_base + (u32)&__bss_end;
	while (bss_start < bss_end) *bss_start++ = 0;

	/* Relocate the GOT. */
	u32* got_start = (u32*)(new_base + (u32)&__got_start);
	u32* got_end = (u32*)(new_base + (u32)&__got_end);
	while (got_start < got_end) *got_start++ += (u32)new_base;

	/* Jump to entry point ("_main") at relocated payload. */
	ep = (ep_t)(new_base + ((u32)_main - (u32)curr_base));
	ep();
}
